# -*- coding: utf-8 -*-
"""
Author: chen xi (modified by Qwen)
Date: 2024/12/24 下午2:30
File: mass_tcp_client.py
功能：高并发 TCP 客户端，发送加密二进制数据，接收时不尝试 UTF-8 解码
"""

import asyncio
import random
import string

from string_enc import iPN_StringEnc


async def connect_and_send_periodically(host, port, interval, message):
    try:
        reader, writer = await asyncio.open_connection(host, port)
        addr = writer.get_extra_info('peername')

        # ✅ 发送加密后的二进制数据（message 已是 bytes 或 str）
        if isinstance(message, str):
            # 如果 message 是字符串（不应如此），先 encode
            send_data = message.encode('utf-8')
        else:
            # 推荐：iPN_StringEnc 返回 bytes
            send_data = message

        async def send_data_task():
            while True:
                await asyncio.sleep(interval)
                try:
                    writer.write(send_data)
                    # print(send_data)
                    await writer.drain()
                except Exception as e:
                    # 发送失败，退出循环
                    print(f"Send error to {addr}: {e}")
                    break

        # 启动发送任务
        asyncio.create_task(send_data_task())

        # 接收循环：只读取数据，不 decode
        try:
            while True:
                data = await reader.read(1024)  # 可根据协议调整大小
                if not data:
                    break

                # ✅ 方法 1：打印为十六进制（推荐，安全）
                print(f"Received {len(data)} bytes from {addr}: {data.hex()}")

                # ✅ 方法 2：打印为 bytes 字面量（可选）
                # print(f"Received {len(data)} bytes from {addr}: {repr(data)}")

                # ✅ 方法 3：如果你想保存原始数据做其他处理
                # process_binary_data(data)

        except Exception as e:
            print(f"Error receiving from {addr}: {type(e).__name__}: {e}")
        finally:
            print(f"Closing connection to {addr}")
            if not writer.is_closing():
                writer.close()
                await writer.wait_closed()

    except Exception as e:
        print(f"Failed to connect to {host}:{port}: {e}")


async def create_connections(server_id, num_connections, host, port, interval, did_end=None):
    tasks = []
    if did_end is None or not did_end:
        did_end = ''.join(random.choices(string.ascii_uppercase, k=6))  # 随机生成6位大写字母
        print(f"Using random DID end: {did_end}")

    for i in range(server_id * num_connections, (server_id + 1) * num_connections):
        # 构造原始 DID 字符串
        did = f"IOTGDD-{i:06d}-{did_end}"
        # 加密 → 返回 (status, encrypted_bytes)
        _, encrypted_message = iPN_StringEnc("RSIOTGWakeupKey@", did, 256)

        # ✅ 确保 encrypted_message 是 bytes 类型
        if isinstance(encrypted_message, str):
            # 如果返回的是 hex 或 base64 字符串，请根据实际调整
            # 例如：encrypted_message = bytes.fromhex(encrypted_message)
            pass  # 根据你的 iPN_StringEnc 实际返回处理

        # 创建连接任务
        task = asyncio.create_task(
            connect_and_send_periodically(host, port, interval, encrypted_message)
        )
        tasks.append(task)

        # 可选：防瞬时连接风暴
        if i % 1000 == 0 and i > 0:
            await asyncio.sleep(0.5)

    # 等待所有连接完成（通常永不结束）
    await asyncio.gather(*tasks)


if __name__ == '__main__':
    num_connections = 60000  # 连接数
    server_id = 0   # 服务器 ID,每台服务器ID均需要不相同
    did_end = None  # 随机生成6位后缀
    host = '47.107.56.10'   # 杭州服务器
    port = 50000
    interval = 180  # 每 3 分钟发送一次

    # 运行客户端
    asyncio.run(create_connections(server_id, num_connections, host, port, interval, did_end))